package com.flame.ai.pdf.service;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.flame.ai.pdf.bean.ContentItem;
import com.flame.ai.pdf.bean.ImageItem;
import com.flame.ai.pdf.bean.TableItem;
import com.flame.ai.pdf.bean.TextItem;
import com.flame.ai.pdf.config.OllamaProperties;
import com.flame.ai.pdf.config.PdfMineruProperties;
import com.flame.ai.pdf.response.ApiResponse;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.RestTemplate;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author Jflame
 * @version 1.0.0
 * @since 2025/8/26
 */
@Service
public class MineruApiService {

    @Resource
    private PdfMineruProperties pdfMineruProperties;

    private static final Logger logger = LoggerFactory.getLogger(MineruApiService.class);
    private static final ObjectMapper objectMapper = new ObjectMapper();

    private RestTemplate restTemplate = new RestTemplate();

    private List<ContentItem> contentItems;
    private Map<String, String> resultImages;
    private Map<String, Object> result;

    public Map<String, Object> call(org.springframework.core.io.Resource file) {
        MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
        body.add("files", file);
        body.add("lang_list", pdfMineruProperties.getLangList());
        body.add("backend", pdfMineruProperties.getBackend());
        body.add("parse_method", pdfMineruProperties.getParseMethod());
        body.add("formula_enable", pdfMineruProperties.getFormulaEnable());
        body.add("table_enable", pdfMineruProperties.getTableEnable());
        body.add("start_page_id", pdfMineruProperties.getStartPageId());
        body.add("end_page_id", pdfMineruProperties.getEndPageId());
        body.add("output_dir", pdfMineruProperties.getOutputDir());

        body.add("return_md", false);
        body.add("return_content_list", true);
        body.add("return_images", true);


        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.MULTIPART_FORM_DATA);
        headers.setAccept(List.of(MediaType.APPLICATION_JSON));
        headers.setConnection("keep-alive");

        HttpEntity<MultiValueMap<String, Object>> request = new HttpEntity<>(body, headers);

        ResponseEntity<ApiResponse> response = restTemplate.postForEntity(pdfMineruProperties.getApiUrl(), request, ApiResponse.class);
        parse(response.getBody());
        return this.result;
    }

    private void parse(ApiResponse response) {
        if (response == null || response.getResults() == null) {
            logger.warn("Response or results is null");
            return;
        }
        result = new HashMap<>();
        contentItems = List.of();
        resultImages = Map.of();
        response.getResults().forEach((key, value) -> {
            try {
                String contentListStr = value.getContentList();
                if (contentListStr != null && !contentListStr.isEmpty()) {
                    contentItems = objectMapper.readValue(
                            contentListStr,
                            objectMapper.getTypeFactory().constructCollectionType(List.class, ContentItem.class)
                    );
                }
                resultImages = value.getImages();
                result.put("texts", new ArrayList<TextItem>());
                result.put("images", new ArrayList<ImageItem>());
                result.put("tables", new ArrayList<TableItem>());
                parseContentList();
            } catch (Exception e) {
                logger.error("Failed to parse content list for entry: {}", key, e);
            }
        });
    }

    private void parseContentList() {
        for (int i = 0; i < contentItems.size(); i++) {
            String type = contentItems.get(i).getType();
            switch (type) {
                case "text":case "equation":
                    parseText(i);
                    break;
                case "image":
                    parseImage(i);
                    break;
                case "table":
                    parseTable(i);
                    break;
                default:
                    break;
            }
        }
    }

    private void parseText(int index) {
        ContentItem item = contentItems.get(index);

        List<TextItem> texts = (List<TextItem>) result.get("texts");
        if (item.getTextLevel() != null && item.getTextLevel() > 0) {
            TextItem textItem = new TextItem();
            textItem.setText(item.getText());
            textItem.setPage(item.getPageIdx());
            texts.add(textItem);
        } else {
            TextItem latest = texts.getLast();
            latest.setText(latest.getText() + item.getText());
            latest.setPage(item.getPageIdx());
            latest.setType("text");
        }

    }

    private void parseImage(int index) {
        ContentItem item = contentItems.get(index);
        String imagePath = item.getImgPath();
        if (imagePath == null) {
            for (int i = index + 1; i < contentItems.size(); i++) {
                ContentItem contentItem = contentItems.get(i);
                if (contentItem.getType().equals("image")) {
                    contentItem.addImageCaption(contentItem.getImageCaption());
                    contentItem.addImageFootnote(contentItem.getImageFootnote());
                    break;
                }
            }
            return;
        }
        String imgKey = imagePath.split("/")[1];

        ImageItem imageItem = new ImageItem();
        imageItem.setImage(resultImages.get(imgKey));
        imageItem.setCaption(item.getImageCaption());
        imageItem.setFootnote(item.getImageFootnote());
        imageItem.setPage(item.getPageIdx());
        imageItem.setType("image");

        String context = "";
        if (index > 0) {
            ContentItem prevItem = contentItems.get(index - 1);
            context = prevItem.getText();
        }
        if (index < contentItems.size() - 1) {
            ContentItem nextItem = contentItems.get(index + 1);
            context += nextItem.getText();
        }
        imageItem.setContext(context);

        List<ImageItem> images = (List<ImageItem>) result.get("images");
        if (images != null) {
            images.add(imageItem);
        }
    }

    private void parseTable(int index) {
        ContentItem item = contentItems.get(index);

        TableItem tableItem = new TableItem();
        tableItem.setBody(item.getTableBody());
        tableItem.setCaption(item.getTableCaption());
        tableItem.setFootnote(item.getTableFootnote());
        tableItem.setPage(item.getPageIdx());
        tableItem.setType("table");

        String context = "";
        if (index > 0) {
            ContentItem prevItem = contentItems.get(index - 1);
            context = prevItem.getText();
        }
        if (index < contentItems.size() - 1) {
            ContentItem nextItem = contentItems.get(index + 1);
            context += nextItem.getText();
        }
        tableItem.setContext(context);

        List<TableItem> tables = (List<TableItem>) result.get("tables");
        if (tables != null) {
            tables.add(tableItem);
        }
    }


}
